import datetime
from functools import wraps

import jwt
from flask import request, current_app

from app.models import User
from app.utils import Response


def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = None

        # 尝试从请求头中获取 token
        if 'Authorization' in request.headers:
            auth_header = request.headers['Authorization']
            parts = auth_header.split()
            if len(parts) == 2 and parts[0].lower() == 'bearer':
                token = parts[1]

        if not token:
            return Response.error(501,'Token is missing!').to_dict()

        try:
            data = jwt.decode(token, current_app.config['SECRET_KEY'], algorithms=['HS256'])
            # 从 token 中获取 user_id，然后加载用户
            current_user = User.query.get(data['user_id'])
        except Exception as e:
            return Response.error(504, 'Token is invalid or expired!').to_dict()

        # 将 current_user 附加到请求上下文中（可以在视图函数中使用）
        kwargs['current_user'] = current_user
        return f(*args, **kwargs)

    return decorated


def get_jwt_token(user_id):

    payload = {
        'user_id': user_id,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(days=365 * 2)  # 2年有效
    }
    token = jwt.encode(payload, current_app.config['SECRET_KEY'], algorithm='HS256')
    return token